######################################################################
##
## Copyright (C) 2006,  Blekinge Institute of Technology
##
## Filename:      Function.py
## Author:        Simon Kagstrom <ska@bth.se>
## Description:   A function
##
## $Id: Function.py 21087 2009-01-31 06:36:19Z ska $
##
######################################################################
import re, os, cgi

from dissy.Config import config
from dissy.Entity import Entity, AddressableEntity
from dissy.Instruction import Instruction
from dissy.JumpStreamHandler import *
from dissy.StrEntity import StrEntity

ADDRESS_REGEXP  = "[0-9,a-f,A-F]+"
ENCODING_REGEXP = ADDRESS_REGEXP + "[ ]"
INSN_REGEXP     = "[0-9,a-z,A-Z,_,\-,\.,\+]+"
INSN_ARGS_REGEXP= "\**[a-z,A-Z,0-9,_,\,,\(,\),\%,\$,\[,\],!,#,\-, ,&,{,},\*,\+]+"

insnRegExp = re.compile("[ ]*(" + ADDRESS_REGEXP + "):[ \t]+((?:" + ENCODING_REGEXP +")*)[ \t]+(" + INSN_REGEXP + ")+[ \t]*(" + INSN_ARGS_REGEXP + ")*")

class Function(AddressableEntity):
    def __init__(self, fileContainer, address, label, size=0):
	AddressableEntity.__init__(self, address = address, endAddress = address + size, baseAddress = fileContainer.baseAddress)
	self.label = label
	self.all = []
	self.insns = []
	self.file = fileContainer

    def addInstruction(self, insn):
	self.insns.append(insn)
	self.all.append(insn)
	if insn.getAddress() > self.endAddress:
	    self.setSize(insn.getAddress() - self.address)

    def getFile(self):
	return self.file

    def addOther(self, other):
	self.all.append(StrEntity(self, other))

    def lookup(self, address):
	last = self.insns[0]
	for insn in self.insns[1:]:
	    if address >= last.getAddress() and address < insn.getAddress():
		return last
	    last = insn
	return None

    def parse(self, try64bitWorkaround=False):
	"""Parse the function."""
	count = 0
	start, end = self.getExtents()
	if try64bitWorkaround:
	    if start & (1<<31):
		start = long(start) | 0xffffffff00000000
	    if end & (1<<31):
		end = long(end) | 0xffffffff00000000
	s = "%s --wide --demangle --source --start-address=0x%Lx --stop-address=0x%Lx %s" % (config.objdump, start, end, self.file.filename)
	self.stream = os.popen(s)
	if self.stream == None:
	    self.prepareParseFunction()
	self.instructions = []
	self.all = []
	firstNonEmpty=False
	for line in self.stream:
	    # Weed away some unneeded stuff
	    if line.startswith("Disassembly of section ") or line.startswith("%s: " % (self.file.filename)):
		continue
	    if not firstNonEmpty and line.strip() == "":
		continue

	    firstNonEmpty=True
	    r = insnRegExp.match(line)
	    if r != None:
		insn = Instruction(self, long("0x" + r.group(1),16), r.group(2), r.group(3), r.group(4))
		self.addInstruction(insn)
		count = count + 1
	    else:
		self.addOther(cgi.escape(line))
	self.stream.close()

	if count == 0 and try64bitWorkaround == False:
	    # If we couldn't add anything interesting, try the 64-bit
	    # workaround (for e.g., MIPS). See http://bugs.debian.org/cgi-bin/bugreport.cgi?bug=369343
	    return self.parse(try64bitWorkaround = True)

    def link(self):
	"""
	Link the local jumps in this function. External jumps are
	linked to the function they reside in.
	"""
	unresolved = []
	for insn in self.insns:
	    if insn.hasLink() and not insn.link():
		unresolved.append(insn)

	positive, negative = self.getJumpDicts()

	# Fill in the jumps, forwards...
	self.fillInJumpStreams(positive, False)
	# ... and backwards
	self.all.reverse()
	self.fillInJumpStreams(negative, True)
	self.all.reverse()

	return unresolved

    def getInstructions(self):
	return self.insns

    def getAll(self):
	return self.all

    def getJumpDicts(self):
	"""
	Get jump dictionaries for the forward and backward jumps.
	"""
	positive = {}
	negative = {}
	for insn in self.insns:
	    other = insn.getOutLink()
	    if isinstance(other, Instruction) and other != insn:
		if other.getAddress() < insn.getAddress():
		    # Jump from insn to other BACKWARDS
		    negative[insn.getAddress()] = (insn, other)
		else:
		    # Jump from insn to other FORWARDS
		    positive[insn.getAddress()] = (insn, other)
	return positive, negative

    def fillInJumpStreams(self, jumpDict, left):
	"""
	Fill in the jump streams for a dictionary of start
	addresses. Specify if the left or right streams should be
	generated.
	"""
	jumpStreamHandler = JumpStreamHandler()
	for insn in self.all:
	    if isinstance(insn, Instruction):
		# Something starts on this address
		if jumpDict.has_key(insn.getAddress()):
		    stream = jumpStreamHandler.alloc()
		    if stream != None:
			stream.start(jumpDict[insn.getAddress()] )
	    jumpStreamHandler.update(insn)
	    if left == True:
		insn.left_state = jumpStreamHandler.getStateTuple()
	    else:
		insn.right_state = jumpStreamHandler.getStateTuple()


    def __str__(self):
	out = str("0x%08x" % self.address) + " " + str(self.label) + ":\n"
	for insn in self.insns:
	    out = out + str(insn) + "\n"
	return out
